from keras import layers, models


class Model:
    def __init__(self, obs_n, act_dim):
        self.act_dim = act_dim
        self.obs_n = obs_n
        self._build_model()

    def _build_model(self):
        hid1_size = 32
        hid2_size = 32
        # ------------------ build evaluate_net ------------------
        model = models.Sequential()
        model.add(layers.Input(shape=(self.obs_n)))
        model.add(layers.Dense(hid1_size, activation="relu", name="l1"))
        model.add(layers.Dense(hid1_size, activation="relu", name="l2"))
        model.add(layers.Dense(self.act_dim, name="l3"))
        model.summary()
        self.model = model
        # ------------------ build target_model ------------------
        target_model = models.Sequential()
        target_model.add(layers.Input(shape=(self.obs_n)))
        target_model.add(layers.Dense(hid2_size, activation="relu", name="l1"))
        target_model.add(layers.Dense(hid2_size, activation="relu", name="l2"))
        target_model.add(layers.Dense(self.act_dim, name="l3"))
        target_model.summary()
        self.target_model = target_model


if __name__ == "__main__":
    model = Model(4, 2)
    print(model.model.summary())
